# Author: Xinwu
# Describe:
# Completion Time:
# Email: lexinwu@outlook.com

if (!require("optparse", quietly = TRUE))
  install.packages("optparse", repo = "https://mirrors.tuna.tsinghua.edu.cn/CRAN/")

usage = "Rscript get_gene_set.R --species 'Homo sapiens' --category 'H' --file_name 'human_hallmark.gmt'"

option_list = list(
  make_option("--species",type = "character",action = "store",default = NULL,help = "物种名称"),
  make_option("--category",type = "character",action = "store",default = NULL,help = "数据集名称"),
  make_option("--subcategory",type = "character",action = "store",default = NULL,help = "子数据集名称"),
  make_option("--file_name",type = "character",action = "store",default = NULL,help = "文件保存名称")
)
opt_parser = OptionParser(option_list = option_list,usage = usage)
opt = parse_args(opt_parser)

# use msigdbr package to get gene_set for GSEA or GSVA analysis
get_gene_set <- function(species = NULL,
                         category = "H",
                         subcategory = NULL){
  library(msigdbr)
  library(dplyr)
  available_species <- msigdbr_species()$species_name
  if(!(species %in% available_species)){
    stop(paste0('msigdbr not have ',species,
                '!\n Please ensure your species is in the following list: \n',
                paste0(available_species,collapse ='\n')))
  }
  species_data <- msigdbr(species = species)
  available_category <- sort(unique(species_data$gs_cat))
  if(!(category %in% available_category)){
    stop(paste0('Your enter category is no include in ',species,
                '!\nPlease ensure that your input is in the following list:\n',
                paste0(available_category,collapse = ',')))
  }
  if(is.null(subcategory)){
    gene_set <- msigdbr(species = species, category = category)
    gene_set <- gene_set %>% split(x=.$gene_symbol,f = .$gs_name)
  }else{
    species_data <- msigdbr(species = species, category = category)
    available_subcate <- sort(unique(species_data$gs_subcat))
    if(!(subcategory %in% available_subcate)){
      stop(paste0('Your enter subcategory ',subcategory,
                  '\n is not in the following list:\n',
                  paste0(available_subcate,collapse = ',')))
    }else{
      gene_set <- msigdbr(species = species, category = category, subcategory = subcategory)
      gene_set <- gene_set %>% split(x=.$gene_symbol,f = .$gs_name)
    }
  }
  return(gene_set)
}

# sava gene_set data to gmt format file
# gmt 格式的文件第一列为 pathway 第二列为 描述 第三列到最后每一列都是一个基因, 该文件用于 gseapy 做 GSEA 富集分析
sava_gmt <- function(gene_set = NULL,
                     file_name = NULL){
  for(i in names(gene_set)){
    line = paste0(i,'\tna\t',paste0(gene_set[[i]],collapse = "\t"))
    write(x = line,file = file_name,append = T)
  }
}


if(FALSE){
  run_GSEA <- function(rdata = NULL,group_by = NULL,
                     sub_group = NULL,target_group = NULL){
  library(Seurat)
  library(dplyr)
  load(rdata)
  Idents(project) <- group_by
  project <- subset(project, idents = sub_group)
  project_genes <- presto::wilcoxauc(project,group_by)
  target_group_genes <- project_genes %>% dplyr::filter(group == target_group) %>% 
    dplyr::arrange(desc(auc)) %>% dplyr::select(feature,auc)
  write.table(x = target_group_genes,file = paste0(target_group,"_gene.rnk"),
              quote = T,col.names = F,row.names = F,sep = "\t")
  # 得到了 gene.rnk 文件后，使用 gseapy prerank -r gene.rnk -g gmt -o outdir 进行GSEA分析即可。
  }
# GSVA 分析
# project_mx 表达矩阵文件, gene_set 为 get_gene_set function 得到的基因集数据
# gvsa_score <- gsva(project_mx, method = 'gsva', gene_set)
}

gene_set <- get_gene_set(species = opt$species, category = opt$category, subcategory = opt$subcategory)
sava_gmt(gene_set = gene_set, file_name = opt$file_name)